package com.tydic.customtablesource;

import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.accumulators.IntCounter;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.table.api.TableSchema;
import org.apache.flink.table.sources.StreamTableSource;
import org.apache.flink.types.Row;

import java.sql.*;
import java.util.concurrent.TimeUnit;

public class CustomTableSource implements StreamTableSource<Row> {

    /**
     * Type information describing the result type.
     */
    private TypeInformation<Row> returnType;

    /**
     * The schema of the table.
     */
    private TableSchema schema;

    public CustomTableSource(TableSchema schema, TypeInformation<Row> returnType) {
        this.schema = schema;
        this.returnType = returnType;
    }

    @Override
    public DataStream<Row> getDataStream(StreamExecutionEnvironment execEnv) {
        JDBCSource jdbcSource = new JDBCSource();
        return execEnv.addSource(jdbcSource).name("Custom JDBC Source").returns(returnType).setParallelism(1);
    }

    @Override
    public TypeInformation<Row> getReturnType() {
        return returnType;
    }

    @Override
    public TableSchema getTableSchema() {
        return schema;
    }

    @Override
    public String explainSource() {
        return "Custom JDBC Source";
    }

    class JDBCSource extends RichSourceFunction<Row> {

        private Connection conn;
        private PreparedStatement ps;
        private IntCounter currentLimit = new IntCounter(0);

        /**
         * Open方法建立连接，这样不用每次invoke的时候都要建立连接和释放连接
         * @param parameters
         * @throws Exception
         */
        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            String driver = "com.mysql.jdbc.Driver";
            String url = "jdbc:mysql://68.61.64.39:3306/test3?useUnicode=true&characterEncoding=utf-8&allowMultiQueries=true&useSSL=false";
            String userName = "root";
            String password = "tydic12345";
            Class.forName(driver);
            conn = DriverManager.getConnection(url, userName, password);
            getRuntimeContext().addAccumulator("currentLimit", this.currentLimit);
        }

        /**
         * DataStream 调用一次run()方法获取数据
         * @param sourceContext
         * @throws Exception
         */
        @Override
        public void run(SourceContext<Row> sourceContext) throws Exception {
            String sql = "select age,address from user01 limit currentLimit,100";
            while (true) {
                String newSql = StringUtils.replace(sql, "currentLimit", this.currentLimit.getLocalValue() + "");
                System.err.println(newSql);
                ps = conn.prepareStatement(newSql);
                ResultSet resultSet = ps.executeQuery();
                while (resultSet.next()) {
                    Row row = new Row(2);
                    row.setField(0, resultSet.getInt("age"));
                    row.setField(1, resultSet.getString("address"));
                    sourceContext.collect(row);
                    this.currentLimit.add(1);
                }
                TimeUnit.SECONDS.sleep(10);
            }
        }

        @Override
        public void cancel() {
            try {
                ps.cancel();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }

        @Override
        public void close() throws Exception {
            super.close();
            conn.close();
        }
    }
}
